rng("shuffle")
%% import wrappers, do not re-run this section
import com.comsol.model.*
import com.comsol.model.util.*

%% control parameters, can rerun this section to update settings
AutoSave = 1; % for continously generating current densities
Iterations = 5000; % total number of random current profiles, keep this number < 9999
runNum = 75489; % the start index of iterations, this affects the lable of generated data
Iterations=Iterations+runNum-1; % while loop step starts from runNum to Iterations+runNum-1, 
exportStepSize = 0.005; % unit in mm
enableFineMesh = 0; % refine boundary of mesh % disable, just physics extreme fine mesh
enableWholeBodyFineMesh = 0; % disable, just physics extreme fine mesh
FineMeshMaxLBThreshold = 0.005; % do not use a mesh size upper bound lower than this value % disable, this number should not be used
widthIncrementalDecay = 0.002; % if width too small, then at sharp corner, the thicken function may not run correctly 
widthLBThreshold = 0.01; % we do not want width smaller than this number, the current density can then be really big
cleanupconstruct = 1;
outofplanethickness = '0.00001'; % out of plane thickness, need a string value
filepathmain = "F:\\Niko_temp\"; % main folder for exporting current data;
%filepathmain = "\\10.229.62.137\data\Experiment folders\ML Inverse Problem\ML Current Inversion Paper (2024)\training and validation data\type2\"
imgpathmain = "F:\\Niko_temp\"; % main folder for exporting meshed image;
if ~exist(filepathmain, 'dir')
   mkdir(filepathmain)
end
if ~exist(imgpathmain, 'dir')
   mkdir(imgpathmain)
end

%% setup comsol environment, do not re-run this section
modelName = "CD_curves";
modelPath="C:\Users\walsworthlab\Desktop\";
[model, modelcreated] = fsetupCOMSOL(modelName, modelPath);

%% iteratively generate images
while runNum <= Iterations
    disp("----------------------------------------------------------------------------------------")
    disp("Current run # "+num2str(runNum)+" End at "+num2str(Iterations))
    %% create current path
    % remove existing geom objects
    continueCreation = 1;
    forcedAbortinMeshLoop = 0;
    while continueCreation
        %'cc'
        geom_objs=model.component('comp1').geom('geom1').feature.tags;
        %'geom_obj'
        exist_ic1 = 0;
        exist_thi1 = 0;
        %size(geom_objs,1)
        if size(geom_objs,1)>0
            for rowi = 1:size(geom_objs)
                if geom_objs(rowi) == java.lang.String("ic1")
                    %geom_objs(rowi)
                    exist_ic1 = 1;
                end
                if geom_objs(rowi) == java.lang.String("thi1")
                    %geom_objs(rowi)
                    exist_thi1 = 1;
                end
            end
        end
        if exist_thi1
            model.component('comp1').geom('geom1').feature.remove('thi1');
        end
        if exist_ic1
            model.component('comp1').geom('geom1').feature.remove('ic1');
        end
        try
            disp("***Randomize current path geometry...***")
            [width_chosen] = fgenerateRandomWidth()
     
            [controlpoint_num, controlpoint_x,controlpoint_y,crossInterestBox,interest_xmin,interest_xmax,interest_ymin,interest_ymax] = fgenerateRandomSplineCoord();
            % note it turns out spline creation in matlab is not the same as COMSOL
            % so need to look at COMSOL model plot and determine if there is path
            % crossing interest box
            while crossInterestBox == 0 
                'CIB'
                [controlpoint_num, controlpoint_x,controlpoint_y,crossInterestBox,interest_xmin,interest_xmax,interest_ymin,interest_ymax] = fgenerateRandomSplineCoord();
            end
            'reached model'
            model.component('comp1').geom('geom1').lengthUnit('mm');
            model.component('comp1').geom('geom1').create('ic1', 'InterpolationCurve');
            model.component('comp1').geom('geom1').feature('ic1').set('table', [controlpoint_x', controlpoint_y']);
            model.component('comp1').geom('geom1').run('ic1');
            model.component('comp1').geom('geom1').create('thi1', 'Thicken2D');
            model.component('comp1').geom('geom1').feature('thi1').selection('input').set({'ic1'});
            model.component('comp1').geom('geom1').feature('thi1').set('totalthick', width_chosen);
            model.component('comp1').geom('geom1').feature('ic1').set('startcond', 'zerocurv');
            model.component('comp1').geom('geom1').feature('ic1').set('endcond', 'zerocurv');
            model.component('comp1').geom('geom1').runPre('fin');
            model.component('comp1').geom('geom1').run;

            % visualize model
            fig1=figure(1);
            mphgeom(model)
            patches=findobj('Type','patch'); % locate the path patch
            patch_vertices = patches(1).Vertices;
            patchcrossingInterestBox = 0; 
            % we do screening again next, 
            % because spline in COMSOL and MATLAB generations are different
            for rowi = 1:size(patch_vertices,1)
                withinInterestX = interest_xmin<patch_vertices(rowi,1) & patch_vertices(rowi,1)<interest_xmax;
                withinInterestY = interest_ymin<patch_vertices(rowi,2) & patch_vertices(rowi,2)<interest_ymax;
                if withinInterestX+withinInterestY == 2
                    patchcrossingInterestBox = 1;
                    break
                end
            end
            % in cases only four vertices is used to define a patch, the previous
            % method will miss this scenario, we now sequentially track patch
            % segments and see if there is crossing box of interest
            % warning: this condition is fragile, only works for four indices
            xdata = patches(1).XData;
            ydata = patches(1).YData;
            if size(patches(1).Vertices,1) == 4
                for Coli = 1:size(patches(1).XData,2)
                    if patchcrossingInterestBox == 1
                        break
                    end
                    for Rowi = 1:size(patches(1).XData,1)
                        if patchcrossingInterestBox == 1
                            break
                        end
                        % identify the next vertice
                        if Rowi == size(patches(1).XData,1)
                            nextColi = Coli+1;
                            nextRowi = 1;
                        else
                            nextColi = Coli;
                            nextRowi = Rowi+1;
                        end
                        if nextColi > size(patches(1).XData,2)
                            % out of boundary
                            continue
                        else
                            outofLB = xdata(Rowi,Coli) < interest_xmin & xdata(nextRowi,nextColi) < interest_xmin; % check if the segment to the left of left boundary
                            outofRB = xdata(Rowi,Coli) > interest_xmax & xdata(nextRowi,nextColi) > interest_xmax; % check if the segment to the right of right boundary
                            outofLoB = ydata(Rowi,Coli) < interest_ymin & ydata(nextRowi,nextColi) < interest_ymin; % check if the segment to the south of lower boundary
                            outofUpB = ydata(Rowi,Coli) > interest_xmax & ydata(nextRowi,nextColi) > interest_ymax; % check if the segment to the north of upper boundary
                            if outofLB+outofRB+outofLoB+outofUpB>=1
                                continue
                            else
                                % there is still a chance that the segment
                                % doesn't cross the boundary
                                % the equation of line segment is 
                                testpts_x = linspace(xdata(Rowi,Coli),xdata(nextRowi,nextColi),500);
                                testpts_y = ((ydata(nextRowi,nextColi)-ydata(Rowi,Coli))/(xdata(nextRowi,nextColi)-xdata(Rowi,Coli))).*(testpts_x-xdata(Rowi,Coli))+ydata(nextRowi,nextColi);
                                inx = interest_xmin<testpts_x & testpts_x<interest_xmax;
                                iny = interest_ymin<testpts_y & testpts_y<interest_ymax;
                                overall = sum(inx&iny);
                                if overall <5 % at least 1% of the line segment should be inside the box
                                    continue
                                else
                                    patchcrossingInterestBox = 1;
                                end
                            end
                        end
                    end
                end
            end
            hold on
            interestbox_xpts = [interest_xmin, interest_xmax,interest_xmax, interest_xmin,interest_xmin];
            interestbox_ypts = [interest_ymin, interest_ymin,interest_ymax, interest_ymax,interest_ymin];
            plot(interestbox_xpts, interestbox_ypts,'r-')
            hold off
            
            %% mesh
            disp("***Create mesh...***")
            creatingMesh = 1;
            while creatingMesh
                model.component('comp1').mesh('mesh1').autoMeshSize(1);
                try
                    model.component('comp1').mesh('mesh1').run;
                    creatingMesh =0;
                    disp("Mesh created")
                    forcedAbortinMeshLoop = 0;
                catch
                        disp("Encounterd error at thickening step... reduce size again")
                        width_chosen_temp = width_chosen - widthIncrementalDecay;
                        if width_chosen_temp < widthLBThreshold
                            disp("death spiral, second step... need to use really small width, not allowed,  manually abort please, could consider delete some vertices")
                            forcedAbortinMeshLoop = 1;
                        else
                            disp("reduce current width...again")
                            width_chosen = width_chosen_temp;
                        end
                        model.component('comp1').geom('geom1').feature('thi1').set('totalthick', width_chosen);
                        model.component('comp1').geom('geom1').run('thi1');
                        model.component('comp1').geom('geom1').run;
                end
            end

            % visualize mesh
            fig2=figure(2);
            mphmesh(model)
            hold on
            plot(interestbox_xpts, interestbox_ypts,'r-')
            title("Mesh "+num2str(runNum,'%04.f'))
            hold off
            
            
            if patchcrossingInterestBox == 0
                warning("No crossing box of interest, regenerate things")
                close(fig1)
            end
            if patchcrossingInterestBox == 1
                if forcedAbortinMeshLoop == 0
                    continueCreation = 0;
                    disp("Valid current path geometry created!")
                end
            end
        catch E
            warning("a problem occured... regenerating things");
            model.component('comp1').geom('geom1').feature.remove('thi1');
            model.component('comp1').geom('geom1').feature.remove('ic1');
        end
    end
    %% import material, no need to change
    if size(model.component('comp1').material.tags,1) == 0
        disp("***defining materials...***")
        model.component('comp1').material.create('mat1', 'Common');
        model.component('comp1').material('mat1').propertyGroup.create('Enu', 'Young''s modulus and Poisson''s ratio');
        model.component('comp1').material('mat1').propertyGroup.create('linzRes', 'Linearized resistivity');
        model.component('comp1').material('mat1').label('Copper');
        model.component('comp1').material('mat1').set('family', 'copper');
        model.component('comp1').material('mat1').propertyGroup('def').set('relpermeability', {'1' '0' '0' '0' '1' '0' '0' '0' '1'});
        model.component('comp1').material('mat1').propertyGroup('def').set('electricconductivity', {'5.998e7[S/m]' '0' '0' '0' '5.998e7[S/m]' '0' '0' '0' '5.998e7[S/m]'});
        model.component('comp1').material('mat1').propertyGroup('def').set('heatcapacity', '385[J/(kg*K)]');
        model.component('comp1').material('mat1').propertyGroup('def').set('relpermittivity', {'1' '0' '0' '0' '1' '0' '0' '0' '1'});
        model.component('comp1').material('mat1').propertyGroup('def').set('emissivity', '0.5');
        model.component('comp1').material('mat1').propertyGroup('def').set('density', '8940[kg/m^3]');
        model.component('comp1').material('mat1').propertyGroup('def').set('thermalconductivity', {'400[W/(m*K)]' '0' '0' '0' '400[W/(m*K)]' '0' '0' '0' '400[W/(m*K)]'});
        model.component('comp1').material('mat1').propertyGroup('Enu').set('E', '');
        model.component('comp1').material('mat1').propertyGroup('Enu').set('nu', '');
        model.component('comp1').material('mat1').propertyGroup('Enu').set('E', '');
        model.component('comp1').material('mat1').propertyGroup('Enu').set('nu', '');
        model.component('comp1').material('mat1').propertyGroup('Enu').set('E', '126e9[Pa]');
        model.component('comp1').material('mat1').propertyGroup('Enu').set('nu', '0.34');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('rho0', '');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('alpha', '');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('Tref', '');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('rho0', '');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('alpha', '');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('Tref', '');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('rho0', '1.667e-8[ohm*m]');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('alpha', '3.862e-3[1/K]');
        model.component('comp1').material('mat1').propertyGroup('linzRes').set('Tref', '293.15[K]');
        model.component('comp1').material('mat1').propertyGroup('linzRes').addInput('temperature');
        model.component('comp1').material('mat1').set('family', 'copper');
        disp("Material defined!")
    else
        disp("Material already defined")
    end

    %% define boundary conditions
    model.component('comp1').physics('ec').prop('d').set('d',outofplanethickness);
    [gnd_label ,term_label] = fgenerateRandomGNDTERMLabel();
    [voltage_chosen] = fgenerateRandomVoltage();
    try
        disp("***Create GND feature...***")
        model.component('comp1').physics('ec').create('gnd1', 'Ground', 1);
        model.component('comp1').physics('ec').feature('gnd1').selection.set([gnd_label]); %#ok<*NBRAK>
        disp("GND feature created!")
    catch
        warning("GND feature exists, only updated label")
        model.component('comp1').physics('ec').feature('gnd1').selection.set([gnd_label]);
    end
    try
        disp("***Create V0 feature...***")
        model.component('comp1').physics('ec').create('term1', 'Terminal', 1);
        model.component('comp1').physics('ec').feature('term1').set('TerminalType', 'Voltage');
        model.component('comp1').physics('ec').feature('term1').selection.set([term_label]);
%         model.component('comp1').physics('ec').feature('term1').set('I0', current_chosen);
        model.component('comp1').physics('ec').feature('term1').set('V0', voltage_chosen);
        disp("V0 feature created!")
    catch
        warning("V0 feature exists, only updated label and current")
        model.component('comp1').physics('ec').feature('term1').selection.set([term_label]);
%         model.component('comp1').physics('ec').feature('term1').set('I0', current_chosen);
        model.component('comp1').physics('ec').feature('term1').set('V0', voltage_chosen);
    end

    

    %% study setup and run, no need to modify, do not re-run this section
    disp("***Set and study...***")
    model.sol.create('sol1');
    model.sol('sol1').study('std1');
    model.study('std1').feature('stat').set('notlistsolnum', 1);
    model.study('std1').feature('stat').set('notsolnum', 'auto');
    model.study('std1').feature('stat').set('listsolnum', 1);
    model.study('std1').feature('stat').set('solnum', 'auto');
    model.sol('sol1').create('st1', 'StudyStep');
    model.sol('sol1').feature('st1').set('study', 'std1');
    model.sol('sol1').feature('st1').set('studystep', 'stat');
    model.sol('sol1').create('v1', 'Variables');
    model.sol('sol1').feature('v1').set('control', 'stat');
    model.sol('sol1').create('s1', 'Stationary');
    model.sol('sol1').feature('s1').create('fc1', 'FullyCoupled');
    model.sol('sol1').feature('s1').feature('fc1').set('linsolver', 'dDef');
    model.sol('sol1').feature('s1').feature.remove('fcDef');
    model.sol('sol1').attach('std1');
    model.result.create('pg1', 'PlotGroup2D');
    model.result('pg1').label('Electric Potential (ec)');
    model.result('pg1').set('frametype', 'spatial');
    model.result('pg1').set('showlegendsmaxmin', true);
    model.result('pg1').set('data', 'dset1');
    model.result('pg1').feature.create('surf1', 'Surface');
    model.result('pg1').feature('surf1').set('showsolutionparams', 'on');
    model.result('pg1').feature('surf1').set('solutionparams', 'parent');
    model.result('pg1').feature('surf1').set('colortable', 'Dipole');
    model.result('pg1').feature('surf1').set('showsolutionparams', 'on');
    model.result('pg1').feature('surf1').set('data', 'parent');
    model.result('pg1').feature.create('str1', 'Streamline');
    model.result('pg1').feature('str1').set('showsolutionparams', 'on');
    model.result('pg1').feature('str1').set('solutionparams', 'parent');
    model.result('pg1').feature('str1').set('expr', {'ec.Ex' 'ec.Ey'});
    model.result('pg1').feature('str1').set('titletype', 'none');
    model.result('pg1').feature('str1').set('posmethod', 'uniform');
    model.result('pg1').feature('str1').set('udist', 0.02);
    model.result('pg1').feature('str1').set('maxlen', 0.4);
    model.result('pg1').feature('str1').set('maxtime', Inf);
    model.result('pg1').feature('str1').set('inheritcolor', false);
    model.result('pg1').feature('str1').set('showsolutionparams', 'on');
    model.result('pg1').feature('str1').set('maxtime', Inf);
    model.result('pg1').feature('str1').set('showsolutionparams', 'on');
    model.result('pg1').feature('str1').set('maxtime', Inf);
    model.result('pg1').feature('str1').set('showsolutionparams', 'on');
    model.result('pg1').feature('str1').set('maxtime', Inf);
    model.result('pg1').feature('str1').set('showsolutionparams', 'on');
    model.result('pg1').feature('str1').set('maxtime', Inf);
    model.result('pg1').feature('str1').set('data', 'parent');
    model.result('pg1').feature('str1').selection.geom('geom1', 1);
    model.result('pg1').feature('str1').selection.set([1 2 3 4]);
    model.result('pg1').feature('str1').set('inheritplot', 'surf1');
    model.result('pg1').feature('str1').feature.create('col1', 'Color');
    model.result('pg1').feature('str1').feature('col1').set('colortable', 'DipoleDark');
    model.result('pg1').feature('str1').feature('col1').set('colorlegend', false);
    model.result('pg1').feature('str1').feature.create('filt1', 'Filter');
    model.result('pg1').feature('str1').feature('filt1').set('expr', '!isScalingSystemDomain');
    model.result.create('pg2', 'PlotGroup2D');
    model.result('pg2').label('Electric Field Norm (ec)');
    model.result('pg2').set('frametype', 'spatial');
    model.result('pg2').set('showlegendsmaxmin', true);
    model.result('pg2').set('data', 'dset1');
    model.result('pg2').feature.create('surf1', 'Surface');
    model.result('pg2').feature('surf1').set('showsolutionparams', 'on');
    model.result('pg2').feature('surf1').set('solutionparams', 'parent');
    model.result('pg2').feature('surf1').set('expr', 'ec.normE');
    model.result('pg2').feature('surf1').set('colortable', 'Prism');
    model.result('pg2').feature('surf1').set('colortabletrans', 'nonlinear');
    model.result('pg2').feature('surf1').set('colorcalibration', -0.8);
    model.result('pg2').feature('surf1').set('showsolutionparams', 'on');
    model.result('pg2').feature('surf1').set('data', 'parent');
    model.result('pg2').feature.create('str1', 'Streamline');
    model.result('pg2').feature('str1').set('showsolutionparams', 'on');
    model.result('pg2').feature('str1').set('solutionparams', 'parent');
    model.result('pg2').feature('str1').set('expr', {'ec.Ex' 'ec.Ey'});
    model.result('pg2').feature('str1').set('titletype', 'none');
    model.result('pg2').feature('str1').set('posmethod', 'uniform');
    model.result('pg2').feature('str1').set('udist', 0.02);
    model.result('pg2').feature('str1').set('maxlen', 0.4);
    model.result('pg2').feature('str1').set('maxtime', Inf);
    model.result('pg2').feature('str1').set('inheritcolor', false);
    model.result('pg2').feature('str1').set('showsolutionparams', 'on');
    model.result('pg2').feature('str1').set('maxtime', Inf);
    model.result('pg2').feature('str1').set('showsolutionparams', 'on');
    model.result('pg2').feature('str1').set('maxtime', Inf);
    model.result('pg2').feature('str1').set('showsolutionparams', 'on');
    model.result('pg2').feature('str1').set('maxtime', Inf);
    model.result('pg2').feature('str1').set('showsolutionparams', 'on');
    model.result('pg2').feature('str1').set('maxtime', Inf);
    model.result('pg2').feature('str1').set('data', 'parent');
    model.result('pg2').feature('str1').selection.geom('geom1', 1);
    model.result('pg2').feature('str1').selection.set([1 2 3 4]);
    model.result('pg2').feature('str1').set('inheritplot', 'surf1');
    model.result('pg2').feature('str1').feature.create('col1', 'Color');
    model.result('pg2').feature('str1').feature('col1').set('expr', 'ec.normE');
    model.result('pg2').feature('str1').feature('col1').set('colortable', 'PrismDark');
    model.result('pg2').feature('str1').feature('col1').set('colorlegend', false);
    model.result('pg2').feature('str1').feature('col1').set('colortabletrans', 'nonlinear');
    model.result('pg2').feature('str1').feature('col1').set('colorcalibration', -0.8);
    model.result('pg2').feature('str1').feature.create('filt1', 'Filter');
    model.result('pg2').feature('str1').feature('filt1').set('expr', '!isScalingSystemDomain');
    model.sol('sol1').runAll;
    disp("Study completed!")

    %% dummy two lines, if mouse click the solution label, this line will be executed automatically
    model.result('pg1').run;
    model.result('pg2').run;

    %% export current density results
    if AutoSave == 0
        % in case need to roll back this simulation data
        answer = questdlg('Keep this simulation data?', ...
        'choices', ...
        'Rollback','Save Data','Save Data');
        % Handle response
        switch answer
            case 'Rollback'
                AutoSave = 0;
            case 'Save Data'
                AutoSave = 1;
        end
    end
    if AutoSave == 1
        disp("***Exporting data...***")
        filepathindex = num2str(runNum,'%05.f');
        filepathappendix = ".txt";
        completefilepath = filepathmain+filepathindex+filepathappendix;
        outputx_start = interest_xmin; % unit in mm
        outputx_end = interest_xmax; % unit in mm
        outputx_stepsize = exportStepSize; % unit in mm
        outputx_string =  "range("+num2str(outputx_start)+","+num2str(outputx_stepsize)+","+num2str(outputx_end)+")";
        outputy_start = interest_ymin; % unit in mm
        outputy_end = interest_ymax; % unit in mm
        outputy_stepsize = exportStepSize; % unit in mm
        outputy_string =  "range("+num2str(outputy_start)+","+num2str(outputy_stepsize)+","+num2str(outputy_end)+")";
        model.result.export.create('data1', 'Data');
        model.result.export('data1').setIndex('expr', 'ec.Jx', 0);
        model.result.export('data1').setIndex('expr', 'ec.Jy', 1);
        model.result.export('data1').set('location', 'grid');
        model.result.export('data1').set('gridx2', outputx_string);
        model.result.export('data1').set('gridy2', outputy_string);
        model.result.export('data1').set('header', false);
        model.result.export('data1').set('fullprec', false);
        model.result.export('data1').set('filename', completefilepath);
        model.result.export('data1').run;
        disp("Data exported to "+completefilepath+"!")
        % imgpathappendix = ".png";
        % completeimgpath = imgpathmain+filepathindex+imgpathappendix;
        % pause(0.1); % some time saveas will fail randmly after 400-500 cycles
        % saveas(fig2,completeimgpath); % save meshed data figure
        runNum = runNum + 1; % move to next simulation round
    end
    %% clean up files
    if cleanupconstruct
        disp("***Cleaning up files...***")
        model.result.remove('pg2');
        model.result.remove('pg1');
        model.result.dataset.remove('dset1');
        model.sol.remove('sol1');
        model.component('comp1').physics('ec').feature.remove('term1');
        model.component('comp1').physics('ec').feature.remove('gnd1');
        model.component('comp1').material.remove('mat1');
        model.component('comp1').geom('geom1').feature.remove('thi1');
        model.component('comp1').geom('geom1').feature.remove('ic1');
        model.component('comp1').geom('geom1').run;
        close(fig1);
        close(fig2);
        disp("Ready for next run!")
        pause(0.1)
    end
end
disp("finished")